package cn.codeartist.springboot.component.dynamic.interceptor;

import cn.codeartist.springboot.component.dynamic.annotation.DB;
import cn.codeartist.springboot.component.dynamic.druid.DynamicDataSourceHolder;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.keygen.SelectKeyGenerator;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import java.util.Properties;

/**
 * 动态数据源拦截器，执行流程：
 * <ol>
 * <li>先取AOP配置的数据源</li>
 * <li>再取Mapper层配置的数据源</li>
 * <li>再执行读写分离</li>
 * </ol>
 *
 * @author 艾江南
 * @date 2019/4/2
 */
@Intercepts({
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})})
public class DynamicDataSourceInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Exception {

        // Service层：AOP 选择的数据源

        if (DynamicDataSourceHolder.getDbType() != null) {
            return invocation.proceed();
        }

        // Mapper层：注解选择的数据源

        Object[] objects = invocation.getArgs();
        MappedStatement mappedStatement = (MappedStatement) objects[0];
        String id = mappedStatement.getId();
        String className = id.substring(0, id.lastIndexOf("."));
        Class<?> clazz = Class.forName(className);
        DB db = clazz.getAnnotation(DB.class);
        if (db != null) {
            try {
                DynamicDataSourceHolder.setDbType(db.value());
                return invocation.proceed();
            } finally {
                DynamicDataSourceHolder.resetDbType();
            }
        }

        // 读写分离

        // 获取事务激活状态
        boolean transactionActive = TransactionSynchronizationManager.isActualTransactionActive();
        // 默认选择主库数据源
        String lookupKey = DynamicDataSourceHolder.DB_MASTER;
        // 没有激活事务
        if (!transactionActive) {
            // 获取查询操作，并且不是查询自增主键
            if (SqlCommandType.SELECT.equals(mappedStatement.getSqlCommandType()) && !id.contains(SelectKeyGenerator.SELECT_KEY_SUFFIX)) {
                BoundSql boundSql = mappedStatement.getSqlSource().getBoundSql(objects[1]);
                String sql = boundSql.getSql().replaceAll("[\\t\\n\\r]", " ");
                // 防止在<select>标签中写DML语句
                String regex = ".*insert\\u0020.*|.*delete\\u0020.*|.*update\\u0020.*";
                if (!sql.matches(regex)) {
                    lookupKey = DynamicDataSourceHolder.DB_SLAVE;
                }
            }
        }
        try {
            DynamicDataSourceHolder.setDbType(lookupKey);
            return invocation.proceed();
        } finally {
            DynamicDataSourceHolder.resetDbType();
        }
    }

    @Override
    public Object plugin(Object target) {
        return target instanceof Executor ? Plugin.wrap(target, this) : target;
    }

    @Override
    public void setProperties(Properties properties) {

    }
}
